#coding:utf-8
from PIL import Image,ImageGrab
import keras as K
from keras.models import load_model
import numpy as np
import pymouse,time,string,random,os
from matplotlib import pyplot
from config import cfg
import argparse, configparser
from ipdb import set_trace

parser = argparse.ArgumentParser()
parser.add_argument('--config', required=True, help='configure file, for example, config_thp_doc.ini')

opt = parser.parse_args()
config = configparser.ConfigParser()
config.read(opt.config)

# lat, log, alt coarse position (x0, y0, x1, y1)
#lat_coarse = (1252, 1017, 1337, 1036)
#log_coarse = (1415, 1017, 1507, 1036)
#alt_coarse = (1581, 1017, 1623, 1036)
lat_coarse = cfg.lat_coarse
log_coarse = cfg.log_coarse
alt_coarse = cfg.alt_coarse

#lat_split = [(3,4,12,18),(12,4,21,18),(30,4,39,18),(39,4,48,18),(48,4,57,18),(57,4,66,18),(66,4,75,18),(75,4,84,18)]
#log_split = [(1,4,10,18),(10,4,19,18),(19,4,28,18),(37,4,46,18),(46,4,55,18),(55,4,64,18),(64,4,73,18),(73,4,82,18),(82,4,91,18)]
#alt_split = [(5,4,14,18),(14,4,23,18),(23,4,32,18),(32,4,41,18)]
lat_split = cfg.lat_split
log_split = cfg.log_split
alt_split = cfg.alt_split

#lat_format = '%s%s.%s%s%s%s%s%s'
#log_format = '%s%s%s.%s%s%s%s%s%s'
#alt_format = '%s%s%s%s'
lat_format = cfg.lat_format
log_format = cfg.log_format
alt_format = cfg.alt_format

#craw_rect = (516, 417, 1582, 661) # dakang -> qiguan
#craw_rect = (373, 165, 1866, 948)
craw_rect = eval(config['craw_info']['craw_rect'])#(239,137,1243,581)
craw_interval =eval(config['craw_info']['craw_interval'])#3

tail_rect = config['craw_info']['tail_rect']
m = pymouse.PyMouse()

def get_coarse_imgs(image, coarse_pos):
    img_coarse = image.crop(coarse_pos)
    img_coarse = img_coarse.convert('L')
    return img_coarse


def split_digits(img_coarse, dig_split):
    digits = []
    for pos in dig_split:
        digits.append(np.array(img_coarse.crop(pos)))
    return digits

def recog_from_digits_imgs(model, digits, dig_format):
    """
    return recognized digit string from digit images
    :param model: keras model
    :param digits: a list of digits images
    :param dig_format: format of this list, like '%s%s.%s%s%s%s%s%s'
    :return: a string, like 83.235564
    """
    x = np.array(digits).astype('float32')
    #x = x.transpose((2,0,1))
    x = x.reshape((x.shape[0], x.shape[1]*x.shape[2]))
    x /= 255
    c = model.predict_classes(x, batch_size=1, verbose=0)
    c = c.tolist()
    return dig_format % tuple(c)


def recog_from_img(model, image, coarse_pos, dig_split, dig_format):
    """
    return digit string
    :param model: keras model
    :param image: whole image
    :param coarse_pos: coarse location of digits
    :param dig_split: split location of every digit
    :param dig_format: digit format, like '%c%c%c.%c%c%c%c'
    :return: digit string
    """
    img_coarse = get_coarse_imgs(image, coarse_pos)
    digits = split_digits(img_coarse, dig_split)
    dig_str = recog_from_digits_imgs(model, digits, dig_format)
    return dig_str, img_coarse


def craw_ge(model):
    """
    The high level api
    :param model: keras model
    :return: lat, log, alt
    """
    res_txt = open('res_%s_%s.txt' % (config['craw_info']['from_'], config['craw_info']['to_']), 'w')
    k = 0
    for i in range(craw_rect[0], craw_rect[2], craw_interval):
        for j in range(craw_rect[1], craw_rect[3], craw_interval):
            m.move(i, j)
            if k == 0:
                time.sleep(1)
            time.sleep(0.03)
            #print(i,j)
            image = ImageGrab.grab()
            dig_str_lat, img_lat = recog_from_img(model, image, lat_coarse, lat_split, lat_format)
            dig_str_lon, img_lon = recog_from_img(model, image, log_coarse, log_split, log_format)
            dig_str_alt, img_alt = recog_from_img(model, image, alt_coarse, alt_split, alt_format)
            print('lat=', dig_str_lat, 'lon=',dig_str_lon, 'alt=', dig_str_alt, 'x=', i, 'y=', j)
            img_id = '%.5d' % k
            res_txt.write('image_id={}, lat={}, lon={}, alt={}, x = {}, y = {}\n'.format(img_id, dig_str_lat, dig_str_lon, dig_str_alt, i-craw_rect[0], j-craw_rect[1]))
            #img_lat = np.array(img_lat)
            #img_lon = np.array(img_lon)
            #img_alt = np.array(img_alt)
            #img_all = Image.fromarray(np.hstack((img_lat, img_log, img_alt)))
            #img_all.save(os.path.join('res_%s_%s' % (cfg.from_, cfg.to_), img_id+'.jpg'))
            k = k + 1
    res_txt.close()

def craw_tail(model, offset):
    """
    craw data for a tail
    offset: offset of tail, in screen pixel
    """
    res_txt = open('res_%s_%s.txt' % (config['craw_info']['from_'], config['craw_info']['to_']), 'wa')
    k = 0
    for i in range(tail_rect[0], tail_rect[2], craw_interval):
        for j in range(tail_rect[1], craw_rect[3], craw_interval):
            m.move(i, j)
            if k == 0:
                time.sleep(1)
            time.sleep(0.03)
            image = ImageGrab.grab()
            
            dig_str_lat, img_lat = recog_from_img(model, image, lat_coarse, lat_split, lat_format)
            dig_str_lon, img_lon = recog_from_img(model, image, log_coarse, log_split, log_format)
            dig_str_alt, img_alt = recog_from_img(model, image, alt_coarse, alt_split, alt_format)
            global_x = i-tail_rect[0]+offset[0]
            global_y = j-tail_rect[1]+offset[1]
            print('lat=', dig_str_lat, 'lon=',dig_str_log, 'alt=', dig_str_alt,
                 'x=', global_x, 'y=', global_y)
            res_txt.write('lat={}, lon={}, alt={}, x={}, y={}\n'.format(dig_str_lat, dig_str_lon, dig_str_alt,
                global_x, global_y))
            k = k + 1
    res_txt.close()

def move_tail_right():
    time.sleep(1)
    m.move(config['craw_info']['tail_rect'][2], config['craw_info']['tail_rect'][1])
    time.sleep(2)
    m.drag(config['craw_info']['tail_rect'][0], config['craw_info']['tail_rect'][1])
    time.sleep(2)
    m.release(config['craw_info']['tail_rect'][0], config['craw_info']['tail_rect'][1])
    time.sleep(2)

def move_tail_down():
    time.sleep(1)
    m.move(config['craw_info']['tail_rect'][0], config['craw_info']['tail_rect'][3])
    time.sleep(2)
    m.drag(config['craw_info']['tail_rect'][0], config['craw_info']['tail_rect'][1])
    time.sleep(2)
    m.release(config['craw_info']['tail_rect'][0], config['craw_info']['tail_rect'][1])
    time.sleep(2)

def move_tail_up():
    time.sleep(1)


#def craw_ge_tail(model):
    #"""
    #craw data from ge
    #"""
    #tail_width = config['craw_info']['tail_rect'][2] - config['craw_info']['tail_rect'][0]
    #tail_height = config['craw_info']['tail_rect'][3] - config['craw_info']['tail_rect'][1]
    #for tail_x in range(config['craw_info']['num_tail'][0]):
        #for tail_y in range(config['craw_info']['num_tail'][1]):
            #tail_offset = (tail_x * tail_width, tail_y * tail_height)
            #craw_tail(model, tail_offset)
            #if tail_y != config['craw_info']['num_tail'][1] - 1:
                #move_tail_vertical()
            #else:



# main

#test craw google earth
print('Pauss for 5 seconds')
print('Open Google Earth in this period')
time.sleep(5)
model = load_model(cfg.ocr_model)
craw_ge(model)

# test by image
# model = load_model(cfg.ocr_model)
# image = Image.open('test.png')
# set_trace()
# dig_str_lat, img_lat = recog_from_img(model, image, lat_coarse, lat_split, lat_format)
# dig_str_lon, img_lon = recog_from_img(model, image, log_coarse, log_split, log_format)
# dig_str_alt, img_alt = recog_from_img(model, image, alt_coarse, alt_split, alt_format)